import pandas as pd
import numpy as np
import argparse

def calc_ht(age, sex, C0, C1, C2):
    return C0 + C1 * age + C2 * age ** 2

def calc_wt(age, sex, C0, C1):
    return C0 - C1 * (age - 45) ** 2

def calc_bsa(wt, ht, c1_param, weight_exp, height_exp):
    return c1_param * (wt ** weight_exp) * (ht ** height_exp)

def generate_china_population(
    n=100,
    min_age=20,
    max_age=75,
    female_ratio=0.5,
    age_dist='Uniform',
    weibull_alpha=2.0,
    weibull_beta=10.0,
    prop_below_65=0.5,
    bsa_c1=0.007184,
    bsa_weight_exp=0.425,
    bsa_height_exp=0.725,
    wt_params=None,
    ht_params=None,
    cv_wt_m=16,
    cv_wt_f=17.5,
    cv_ht_m=3.7,
    cv_ht_f=3.7,
    seed=42,
    output_csv='china_population_demo.csv'
):
    np.random.seed(seed)
    n_female = int(n * female_ratio)
    n_male = n - n_female
    # 年龄分布
    if age_dist == 'Uniform':
        ages_male = np.random.uniform(min_age, max_age, n_male)
        ages_female = np.random.uniform(min_age, max_age, n_female)
    elif age_dist == 'Weibull':
        ages_male = weibull_beta * np.random.weibull(weibull_alpha, n_male) + min_age
        ages_female = weibull_beta * np.random.weibull(weibull_alpha, n_female) + min_age
        ages_male = np.clip(ages_male, min_age, max_age)
        ages_female = np.clip(ages_female, min_age, max_age)
    elif age_dist == 'Split':
        n_male_below = int(n_male * prop_below_65)
        n_male_above = n_male - n_male_below
        n_female_below = int(n_female * prop_below_65)
        n_female_above = n_female - n_female_below
        ages_male = np.concatenate([
            np.random.uniform(min_age, 65, n_male_below),
            weibull_beta * np.random.weibull(weibull_alpha, n_male_above) + 65
        ])
        ages_female = np.concatenate([
            np.random.uniform(min_age, 65, n_female_below),
            weibull_beta * np.random.weibull(weibull_alpha, n_female_above) + 65
        ])
        ages_male = np.clip(ages_male, min_age, max_age)
        ages_female = np.clip(ages_female, min_age, max_age)
    else:
        raise ValueError('age_dist must be Uniform, Weibull, or Split')
    # 参数
    if wt_params is None:
        wt_params = {'M': (75, 0.005888), 'F': (65, 0.010767)}
    if ht_params is None:
        ht_params = {'M': (175.32, 0.1113, -0.0025), 'F': (161.66, 0.1319, -0.0027)}
    # 身高、体重
    heights_male = calc_ht(ages_male, 'M', *ht_params['M'])
    heights_female = calc_ht(ages_female, 'F', *ht_params['F'])
    weights_male = calc_wt(ages_male, 'M', *wt_params['M'])
    weights_female = calc_wt(ages_female, 'F', *wt_params['F'])
    # 应用CV
    heights_male *= np.random.normal(1, cv_ht_m / 100, n_male)
    heights_female *= np.random.normal(1, cv_ht_f / 100, n_female)
    weights_male *= np.random.normal(1, cv_wt_m / 100, n_male)
    weights_female *= np.random.normal(1, cv_wt_f / 100, n_female)
    # BSA
    bsa_male = calc_bsa(weights_male, heights_male, bsa_c1, bsa_weight_exp, bsa_height_exp)
    bsa_female = calc_bsa(weights_female, heights_female, bsa_c1, bsa_weight_exp, bsa_height_exp)
    # 合并
    df_male = pd.DataFrame({
        'id': np.arange(1, n_male + 1),
        'sex': 'M',
        'age': ages_male,
        'height': heights_male,
        'weight': weights_male,
        'bsa': bsa_male
    })
    df_female = pd.DataFrame({
        'id': np.arange(n_male + 1, n + 1),
        'sex': 'F',
        'age': ages_female,
        'height': heights_female,
        'weight': weights_female,
        'bsa': bsa_female
    })
    df = pd.concat([df_male, df_female], ignore_index=True)
    df.to_csv(output_csv, index=False)
    print(f'已保存到 {output_csv}')
    print(df.head())
    return df

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--n', type=int, default=100)
    parser.add_argument('--min_age', type=int, default=20)
    parser.add_argument('--max_age', type=int, default=75)
    parser.add_argument('--female_ratio', type=float, default=0.5)
    parser.add_argument('--age_dist', type=str, default='Uniform', choices=['Uniform', 'Weibull', 'Split'])
    parser.add_argument('--weibull_alpha', type=float, default=2.0)
    parser.add_argument('--weibull_beta', type=float, default=10.0)
    parser.add_argument('--prop_below_65', type=float, default=0.5)
    parser.add_argument('--bsa_c1', type=float, default=0.007184)
    parser.add_argument('--bsa_weight_exp', type=float, default=0.425)
    parser.add_argument('--bsa_height_exp', type=float, default=0.725)
    parser.add_argument('--cv_wt_m', type=float, default=16)
    parser.add_argument('--cv_wt_f', type=float, default=17.5)
    parser.add_argument('--cv_ht_m', type=float, default=3.7)
    parser.add_argument('--cv_ht_f', type=float, default=3.7)
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--output_csv', type=str, default='china_population_demo.csv')
    args = parser.parse_args()
    generate_china_population(
        n=args.n,
        min_age=args.min_age,
        max_age=args.max_age,
        female_ratio=args.female_ratio,
        age_dist=args.age_dist,
        weibull_alpha=args.weibull_alpha,
        weibull_beta=args.weibull_beta,
        prop_below_65=args.prop_below_65,
        bsa_c1=args.bsa_c1,
        bsa_weight_exp=args.bsa_weight_exp,
        bsa_height_exp=args.bsa_height_exp,
        cv_wt_m=args.cv_wt_m,
        cv_wt_f=args.cv_wt_f,
        cv_ht_m=args.cv_ht_m,
        cv_ht_f=args.cv_ht_f,
        seed=args.seed,
        output_csv=args.output_csv
    ) 